import math
import time
import json
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401
from scipy import ndimage
import pickle
import os
from shapely import wkt
from shapely.geometry import Point, Polygon, MultiPolygon
from coord_transform import PositionConvert

class MapGenerator3D:
    """
    高效的三维栅格与 ESDF 生成器。

    两种建图方式：
      - build_from_box(...) : 直接在栅格层面把外部边界（箱体）和若干内部面/墙设为占据（非常快）
      - build_from_point_cloud(ox, oy, oz) : 将稀疏/任意点云索引到栅格，然后做膨胀（适用于真实点云）

    最后使用 ndimage.distance_transform_edt 计算 ESDF，并用距离阈值执行 robot 半径膨胀。
    """

    def __init__(self, resolution: float, robot_radius: float):
        """
        resolution : 栅格分辨率（米）
        robot_radius : 机器人半径（米），用于膨胀障碍
        """
        self.resolution = float(resolution)
        self.robot_radius = float(robot_radius)
        # 地图的地理参考点
        self.ref_lat = 0.0
        self.ref_lon = 0.0
        # 世界坐标边界（在 build 时计算）
        self.min_x = self.min_y = self.min_z = 0.0
        self.max_x = self.max_y = self.max_z = 0.0

        # 栅格维度（在 build 时计算）
        self.x_width = self.y_width = self.z_width = 0

        # 数据（在 build 时填充）
        self.obs_map = None   # bool ndarray: True 表示占据
        self.esdf_map = None  # float ndarray: 每个栅格到最近障碍的欧式距离 (m)

    # ---------------- 辅助转换 ----------------
    def calc_grid_position(self, index: int, min_position: float) -> float:
        """栅格索引 -> 世界坐标（格子的中心）"""
        return index * self.resolution + min_position

    def calc_xyz_index(self, pos: float, min_pos: float) -> int:
        """世界坐标 -> 栅格索引（round）"""
        return int(round((pos - min_pos) / self.resolution))

    # ---------------- 建图方法 ----------------
    def _setup_bounds(self, min_x, max_x, min_y, max_y, min_z, max_z):
        """内部：计算边界与栅格尺寸"""
        self.min_x = float(math.floor(min_x))
        self.min_y = float(math.floor(min_y))
        self.min_z = float(math.floor(min_z))
        self.max_x = float(math.ceil(max_x))
        self.max_y = float(math.ceil(max_y))
        self.max_z = float(math.ceil(max_z))

        self.x_width = max(1, int(math.ceil((self.max_x - self.min_x) / self.resolution)))
        self.y_width = max(1, int(math.ceil((self.max_y - self.min_y) / self.resolution)))
        self.z_width = max(1, int(math.ceil((self.max_z - self.min_z) / self.resolution)))

        # 初始化占据格子数组
        self.obs_map = np.zeros((self.x_width, self.y_width, self.z_width), dtype=bool)
        self.esdf_map = None

    def _finalize_obs_and_esdf(self):
        """
        给定已标记的初始占据 self.obs_map（True = 原障碍或边界），
        通过 distance transform 实现 robot_radius 膨胀，生成最终 obs_map 与 esdf_map。
        """
        if self.obs_map is None:
            raise RuntimeError("请先初始化 self.obs_map")

        # 计算自由空间的距离到最近占据体素（单位：格子数），乘以 resolution 得到米
        free_space = np.logical_not(self.obs_map)
        # distance_transform_edt 的返回值是以像素/格子为单位的距离
        dist_grid = ndimage.distance_transform_edt(free_space) * self.resolution

        # 将距离小于等于 robot_radius 的格子标为占据（膨胀）
        inflated = dist_grid <= self.robot_radius
        self.obs_map = np.array(inflated, dtype=bool)

        # 最终 ESDF（对膨胀后的 free space 计算距离）
        free_after = np.logical_not(self.obs_map)
        self.esdf_map = ndimage.distance_transform_edt(free_after) * self.resolution

    # ---------------- 保存与可视化 ----------------
    def save_map(self, filename="3d_map_data.pkl"):
        """保存地图数据到文件（pickle）"""
        if self.obs_map is None or self.esdf_map is None:
            raise ValueError("请先构建地图再保存")

        map_data = {
            'ref_lat': self.ref_lat,
            'ref_lon': self.ref_lon,
            'resolution': self.resolution,
            'robot_radius': self.robot_radius,
            'min_x': self.min_x,
            'min_y': self.min_y,
            'min_z': self.min_z,
            'max_x': self.max_x,
            'max_y': self.max_y,
            'max_z': self.max_z,
            'x_width': self.x_width,
            'y_width': self.y_width,
            'z_width': self.z_width,
            'obs_map': self.obs_map,
            'esdf_map': self.esdf_map
        }
        with open(filename, 'wb') as f:
            pickle.dump(map_data, f)
        print(f"[save_map] 地图已保存到: {filename}")
        return filename

    def visualize_map(self, save_path=None, show=True, subsample=1):
        """3D 可视化占据栅格（可选保存）。subsample 用于减少绘图点数（加速）"""
        if self.obs_map is None:
            raise ValueError("请先构建地图再可视化")

        xs, ys, zs = [], [], []
        # 采样输出避免绘图瓶颈
        for ix in range(0, self.x_width, subsample):
            for iy in range(0, self.y_width, subsample):
                for iz in range(0, self.z_width, subsample):
                    if self.obs_map[ix, iy, iz]:
                        xs.append(self.calc_grid_position(ix, self.min_x))
                        ys.append(self.calc_grid_position(iy, self.min_y))
                        zs.append(self.calc_grid_position(iz, self.min_z))

        # 确保创建新图形前，清除当前可能残留的绘图状态
        plt.clf()
        fig = plt.figure(figsize=(8, 8))
        ax = fig.add_subplot(111, projection='3d')
        if xs:
            ax.scatter(xs, ys, zs, c='k', s=6, alpha=0.6, label='obstacle voxels')
        ax.set_box_aspect((max(1, self.x_width), max(1, self.y_width), max(1, self.z_width)))
        ax.set_title('3D Occupancy Grid (voxels)')
        ax.set_xlabel('X [m]'); ax.set_ylabel('Y [m]'); ax.set_zlabel('Z [m]')
        ax.legend()
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"[visualize_map] 可视化已保存到 {save_path}")
        if show:
            plt.show()
        return fig

def build_map_from_environmental_info(
    map_resolution: float = 20.0,
    robot_radius: float = 10.0,
    output_filename: str = "map.pkl",
    env_info: dict = None,
    mission_bounds: dict = None,
    ref_lat: float = None,
    ref_lon: float = None,
) -> MapGenerator3D:
    """
    解析北航请求中的环境信息, 生成三维障碍物点云,
    并使用 MapGenerator3D 创建和保存地图。

    Args:
        beihang_request: 完整的JSON请求 (以Python字典形式)。
        map_resolution: 输出地图的分辨率 (米)。
        robot_radius: 机器人半径, 用于障碍物膨胀 (米)。
        output_filename: 最终地图 pickle 文件的保存路径。
    """
    
    max_map_size_m = 100000.0  #  最大地图尺寸限制 (米)
    padding = 1000.0  # 在最终边界外增加的缓冲 (米)
    print(f"--- 开始从环境信息生成地图,最大尺寸限制为 {max_map_size_m} 米 ---")
    pos_converter = PositionConvert() # 你的坐标转换器
    print(f"使用参考坐标: Lat={ref_lat}, Lon={ref_lon}")
    # 1. 提取任务起点和终点        # 垂直方向
    z_min = 0.0  
    z_max = 1000.0 
    print(f"障碍物垂直范围: Z = {z_min}m ~ {z_max}m")

    # 经纬度边界（解析阶段已经算好）
    if not mission_bounds:
        raise ValueError("缺少 mission_bounds，无法确定地图边界")

    min_lon = mission_bounds["min_lon"]
    max_lon = mission_bounds["max_lon"]
    min_lat = mission_bounds["min_lat"]
    max_lat = mission_bounds["max_lat"]
    # 2. 收集并解析所有的障碍物几何体
    all_obstacle_polygons_xy = []
    
    # 定义障碍物源以及如何从中提取WKT字符串的 lambda 函数
    obstacle_sources = {
        "obstacles": lambda item: item,
        "noFlyZones": lambda item: item,
        "sensitiveAreas": lambda item: item.get("geom")
    }

    for key, extractor in obstacle_sources.items():
        print(f"正在处理障碍物源: '{key}'...")
        count = 0
        for item in env_info.get(key, []):
            wkt_string = extractor(item)
            if not wkt_string:
                continue
            
            try:
                # 解析 WKT 字符串为 shapely 几何对象
                geom = wkt.loads(wkt_string)
                # 统一处理 MultiPolygon 和 Polygon 的情况
                geoms_to_process = geom.geoms if geom.geom_type.startswith('Multi') else [geom]
                
                for poly in geoms_to_process:
                    if poly.geom_type != 'Polygon' or not poly.exterior:
                        continue
                    
                    # 将多边形顶点从 经纬度 转换为 本地XY坐标(米)
                    xy_coords = [pos_converter.GPStoXY(lat, lon, ref_lat, ref_lon) for lon, lat in poly.exterior.coords]
                    all_obstacle_polygons_xy.append(Polygon(xy_coords))
                    count += 1
            except Exception as e:
                print(f"  - 警告: 解析 '{key}' 中的几何体失败。错误: {e}")
        print(f"  - 找到并转换了 {count} 个多边形。")
    
    # 3. 计算边界并检查地图尺寸,计算包含任务点和障碍物的统一边界
    # 我们用四个角来估计 XY 区域
    corner_ll = [
        (min_lat, min_lon),
        (min_lat, max_lon),
        (max_lat, min_lon),
        (max_lat, max_lon),
    ]
    xs = []
    ys = []
    for lat, lon in corner_ll:
        x, y = pos_converter.GPStoXY(lat, lon, ref_lat, ref_lon)
        xs.append(x)
        ys.append(y)
    bounds_min_x = min(xs)
    bounds_max_x = max(xs)
    bounds_min_y = min(ys)
    bounds_max_y = max(ys)

    # 如果有障碍物, 则将障碍物区域合并到边界中
    if all_obstacle_polygons_xy:
        obs_bounds = MultiPolygon(all_obstacle_polygons_xy).bounds
        bounds_min_x = min(bounds_min_x, obs_bounds[0])
        bounds_min_y = min(bounds_min_y, obs_bounds[1])
        bounds_max_x = max(bounds_max_x, obs_bounds[2])
        bounds_max_y = max(bounds_max_y, obs_bounds[3])
    
    map_width = bounds_max_x - bounds_min_x
    map_height = bounds_max_y - bounds_min_y
    print(f"计算出的统一地图边界: X跨度={map_width:.1f}m, Y跨度={map_height:.1f}m")

    # 检查地图规模是否超出限制
    if map_width > max_map_size_m or map_height > max_map_size_m:
        error_msg = f"地图规模过大 (长:{map_width:.1f}m, 宽:{map_height:.1f}m)，超过{max_map_size_m}m的限制，已终止建图。"
        print(f"错误: {error_msg}")
        raise ValueError(error_msg)
    
    # 4. 仅光栅化边界并拉伸为三维点云墙
    print("\n正在将二维多边形边界光栅化为三维点云墙...")
    obstacle_points_3d = []
    z_range = np.arange(z_min, z_max, map_resolution)

    for poly in all_obstacle_polygons_xy:
        boundary_coords = list(poly.exterior.coords)
        for i in range(len(boundary_coords) - 1):
            p1 = boundary_coords[i]
            p2 = boundary_coords[i+1]
            x1, y1 = p1
            x2, y2 = p2
            
            # 计算线段长度并确定插值点数
            dist = math.hypot(x2 - x1, y2 - y1)
            num_steps = max(2, int(math.ceil(dist / map_resolution)))
            
            # 沿线段插值生成点
            interp_x = np.linspace(x1, x2, num_steps)
            interp_y = np.linspace(y1, y2, num_steps)
            
            # 将二维边界点拉伸为三维墙
            for ix, iy in zip(interp_x, interp_y):
                for z in z_range:
                    obstacle_points_3d.append((ix, iy, z))

    print(f"总计生成了 {len(obstacle_points_3d)} 个三维障碍物点。")
    
    # 5. 构建并保存地图
    print("\n初始化 MapGenerator3D 并构建地图...")
    map_generator = MapGenerator3D(resolution=map_resolution, robot_radius=robot_radius)
    map_generator.ref_lat = ref_lat
    map_generator.ref_lon = ref_lon
    
    # 步骤 A: 使用最终计算出的边界（加上padding）创建空的栅格地图
    map_generator._setup_bounds(
        bounds_min_x - padding, bounds_max_x + padding,
        bounds_min_y - padding, bounds_max_y + padding,
        z_min, z_max
    )
    
    # 步骤 B: 如果有障碍物点, 将它们“绘制”到已创建的栅格地图上
    if obstacle_points_3d:
        ox, oy, oz = zip(*obstacle_points_3d)
        ox_a, oy_a, oz_a = np.asarray(ox), np.asarray(oy), np.asarray(oz)
        
        # 将点云坐标转换为栅格索引
        xi = np.round((ox_a - map_generator.min_x) / map_resolution).astype(int)
        yi = np.round((oy_a - map_generator.min_y) / map_resolution).astype(int)
        zi = np.round((oz_a - map_generator.min_z) / map_resolution).astype(int)

        # 保证索引在有效范围内
        xi = np.clip(xi, 0, map_generator.x_width - 1)
        yi = np.clip(yi, 0, map_generator.y_width - 1)
        zi = np.clip(zi, 0, map_generator.z_width - 1)
        
        # 标记占据栅格
        map_generator.obs_map[xi, yi, zi] = True

    # 步骤 C: 最终进行膨胀和ESDF计算
    map_generator._finalize_obs_and_esdf()

    print(f"地图生成完毕。栅格维度: {map_generator.x_width}x{map_generator.y_width}x{map_generator.z_width}")
    map_generator.save_map(output_filename)
    print(f"--- 地图已成功保存至 {output_filename} ---")

    return map_generator


